sub_sp = '''
    @SP
    AM=M-1'''

add_sp = '''
    @SP
    M=M+1'''

operand_symbol_code = {
    'add': 'D+M',
    'sub': 'M-D',
    'neg': '-D',
    'eq': 'JEQ',
    'gt': 'JGT',
    'lt': 'JLT',
    'and': 'M&D',
    'or': 'M|D',
    'not': '!D',
}

PushCmd = {
    "constant":
    '''
    @{}
    D=A
    @SP 
    A=M
    M=D'''+add_sp,

    "pointer":
    '''
    @{}
    D=M
    @SP
    A=M
    M=D'''+add_sp,

    "static":
    '''
    @{}
    D=M
    @SP
    A=M
    M=D'''+add_sp,

    "normal":
    '''
    @{}
    D=A
    @{}
    A=D+M
    D=M
    @SP
    A=M
    M=D'''+add_sp,
}

PopCmd = {
    "normal":
    '''
    @{}
    D=A
    @{}
    A=D+M
    D=A
    @R13
    M=D'''+sub_sp +
    '''
    D=M
    @R13
    A=M
    M=D''',

    "pointer":
    sub_sp +
    '''
    D=M
    @{}
    M=D''',

    "static":
    sub_sp +
    '''
    D=M
    @{}
    M=D'''
}

normal_segment = {
    'local': 'LCL',
    'argument': 'ARG',
    'this': 'THIS',
    'that': 'THAT',
}


def get_new_label(index) -> str:
    return "label" + str(index)


def two_operand_code(symbol) -> str:
    symbol = operand_symbol_code[symbol]
    return sub_sp + \
        '''
    D=M''' +\
        sub_sp +\
        f'''
    M={symbol}''' +\
        add_sp


def one_operand_code(symbol) -> str:
    symbol = operand_symbol_code[symbol]
    return sub_sp + \
        '''
    D=M''' +\
        f'''
    @SP
    A=M
    M={symbol}''' + add_sp


def jmp_code(symbol, label) -> str:
    symbol = operand_symbol_code[symbol]
    return sub_sp + \
        '''
    D=M''' + sub_sp +\
        f'''
    D=M-D
    M=-1
    @{label}
    D;{symbol}
    @SP
    A=M
    M=0
({label})''' + add_sp


def push_code(srcfile, segment, offset) -> str:
    if segment in ['pointer', 'temp']:
        reg = ""
        if segment == 'pointer':
            if offset == '0':
                reg = "THIS"
            elif offset == '1':
                reg = "THAT"
        else:
            # temp 0~7 => R 5~12
            reg = "R"+str(int(offset) + 5)
        return PushCmd['pointer'].format(reg)
    elif segment == 'constant':
        return PushCmd['constant'].format(offset)
    elif segment == 'static':
        return PushCmd['static'].format(srcfile + "." + offset)
    else:
        segment = normal_segment[segment]
        return PushCmd['normal'].format(offset, segment)


def pop_code(srcfile, segment, offset) -> str:
    if segment in ['pointer', 'temp']:
        reg = ""
        if segment == 'pointer':
            if offset == '0':
                reg = "THIS"
            elif offset == '1':
                reg = "THAT"
        else:
            # temp 0~7 => R 5~12
            reg = "R"+str(int(offset) + 5)
        return PopCmd['pointer'].format(reg)
    elif segment == 'static':
        return PopCmd['static'].format(srcfile + "." + offset)
    else:
        segment = normal_segment[segment]
        return PopCmd['normal'].format(offset, segment)


def label_code(label) -> str:
    return f'\n({label})'


def if_goto_code(label) -> str:
    return sub_sp + \
        f'''
    D=M
    @{label}
    D;JNE'''


def goto_code(label) -> str:
    return f'''
    @{label}
    0;JMP'''

# SP=256
# call Sys.init 0


def init_code() -> str:
    return '''
    @256
    D=A
    @SP
    M=D''' + call_code("Sys.init", '0')
    #  @Sys.init
    #  0;JMP'''


def function_code(functionName, numLocals) -> str:
    push0 = '''
    @SP
    A=M
    M=0
    '''+add_sp
    return label_code(functionName) + push0 * int(numLocals)


# FRAME = LCL
# RET = *(FRAME-5), IF there are no argument, ARG == FRAME - 5!!!
# *ARG = pop()
# SP = ARG+1
# THAT = *(FRAME-1)
# THIS = *(FRAME-2)
# ARG = *(FRAME-3)
# LCL = *(FRAME-4)
# goto RET
def return_code() -> str:
    pop_code_r14 = '''
    @R14
    AM=M-1
    D=M
    @{}
    M=D'''

    return '''
    @LCL
    D=M
    @R14
    M=D
    @5
    D=A
    @R14
    A=M-D
    D=M
    @R15
    M=D
    ''' + \
        '''
    @SP
    AM=M-1
    D=M
    @ARG
    A=M
    M=D''' + \
        '''
    @ARG
    D=M+1
    @SP
    M=D''' + pop_code_r14.format("THAT") + \
        pop_code_r14.format("THIS") + \
        pop_code_r14.format("ARG") + \
        pop_code_r14.format("LCL") + \
        '''
    @R15
    A=M
    0;JMP'''


call_index = 0

# push return_address
# push LCL
# push ARG
# push THIS
# push THAT
# ARG = SP-n-5
# LCL = SP
# goto f
# (return_address)


def call_code(functionName, numArgs) -> str:
    global call_index
    call_index += 1
    arg_offset = str(int(numArgs) + 5)
    push_code = '''
    @{}
    D=M
    @SP
    A=M
    M=D''' + add_sp
    return f'''
    @RET_ADDRESS_CALL{call_index}
    D=A
    @SP
    A=M
    M=D''' + add_sp +\
        push_code.format('LCL') +\
        push_code.format('ARG') +\
        push_code.format('THIS') +\
        push_code.format('THAT') +\
        f'''
    @{arg_offset}
    D=A
    @SP
    D=M-D
    @ARG
    M=D
    @SP
    D=M
    @LCL
    M=D
    @{functionName}
    0;JMP
(RET_ADDRESS_CALL{call_index})'''
